package com.demo.impl;

import com.NetSDKDemo.ClientDemo;
import org.bytedeco.javacv.FFmpegFrameGrabber;
import org.bytedeco.javacv.FFmpegFrameRecorder;
import org.bytedeco.javacv.Frame;
import org.springframework.stereotype.Controller;

import javax.websocket.OnClose;
import javax.websocket.OnMessage;
import javax.websocket.OnOpen;
import javax.websocket.Session;
import javax.websocket.server.ServerEndpoint;
import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.nio.ByteBuffer;


@ServerEndpoint("/live")
@Controller
public class Websocket {
    /**
     * 与某个客户端的连接会话，需要通过它来给客户端发送数据
     */
    public static Session session;

    private static FFmpegFrameRecorder recorder;
    private static ByteArrayOutputStream outputStream;
    private static boolean initialized = false;

    /**
     * 连接成功
     *
     * @param session
     */
    @OnOpen
    public void onOpen(Session session) throws IOException {
        Websocket.session = session; // 保存客户端连接的Session对象

        outputStream = new ByteArrayOutputStream();
        recorder = new FFmpegFrameRecorder(outputStream, 0);

        ClientDemo.start();
    }

    /**
     * 连接关闭
     *
     * @param session
     */
    @OnClose
    public void onClose(Session session) {


    }

    /**
     * 接收到消息
     *
     * @param text
     */
    @OnMessage
    public String onMsg(String text) throws IOException {
        System.out.println("连接成功");
        return null;
    }

    public static void sendBuffer(byte[] bytes) {
        try {
            // 使用ByteArrayInputStream作为输入流
            ByteArrayInputStream inputStream = new ByteArrayInputStream(bytes);

            // 创建FFmpegFrameGrabber
            FFmpegFrameGrabber grabber = new FFmpegFrameGrabber(inputStream);
            grabber.setFormat("h264");
            grabber.start();

            if (!initialized) {
                initialized = true;
                recorder = new FFmpegFrameRecorder(outputStream, 0);
                recorder.setVideoCodec(grabber.getVideoCodec());
                recorder.setFormat("flv");
                recorder.setFrameRate(grabber.getFrameRate());
                recorder.setGopSize((int) (grabber.getFrameRate() * 2));
                recorder.setVideoBitrate(grabber.getVideoBitrate());
                recorder.setImageWidth(grabber.getImageWidth());
                recorder.setImageHeight(grabber.getImageHeight());
                recorder.start();
            }

            Frame frame;
            while ((frame = grabber.grab()) != null) {
                recorder.record(frame);
            }

            grabber.stop();
            grabber.release();

            byte[] flvData = outputStream.toByteArray();
            System.out.println("flvData size: " + flvData.length);
            outputStream.reset();

            synchronized (session) {
                session.getBasicRemote().sendBinary(ByteBuffer.wrap(flvData));
            }
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }
}
